#ifndef YOLOV8_H_
#define YOLOV8_H_

#include <memory>
#include <deque>
#include <mutex>
#include <thread>
#include <cassert>
#include <atomic>

#define USE_OPENCV 1

#include "bmcv_api.h"
#include "bmcv_api_ext.h"
#include "bmlib_runtime.h"
#include "datapipe.h"
#include "common.h"
#include "opencv2/opencv.hpp"
#include "bmnn_utils.h"
#include "bmlib_runtime.h"
#include "utils.hpp"

#define DEBUG 0

#define MAX_CHANNEL_NUM 32

struct ImageInfo
{
    cv::Size raw_size;
    cv::Vec4d trans;
};

struct Paras
{
    int r_x;
    int r_y;
    int r_w;
    int r_h;
    int width;
    int height;
};

using YoloV8BoxVec = std::vector<YoloV8Box>;

class YoloV8
{

public:
    YoloV8(int dev_id, std::string bmodel_file, std::string tpu_kernel_module_path, std::string mask_bmodel_path, int que_size = 1,
           int skip_num = 0, float nmsThreshold = 0.5, float confThreshold = 0.5, int class_name_num = 6)
        : que_size_(que_size), pre_forward_que_(que_size), forward_post_que_(que_size), out_que_(que_size), skip_num_(skip_num),
          m_nmsThreshold(nmsThreshold), m_confThreshold(confThreshold), m_class_num(class_name_num)
    {

        channel_skip_nums_ = std::vector<int>(MAX_CHANNEL_NUM, 0);
        init(dev_id, bmodel_file, tpu_kernel_module_path, mask_bmodel_path);
        tpumask_Init(mask_bmodel_path, dev_id);
        start();
    }
    ~YoloV8()
    {
        stop();
        release();
    }
    void start();
    void stop();
    void setClassNames(std::vector<std::string> class_name_vec);
    int push_img(int channel_id, std::shared_ptr<bm_image> img);
    std::shared_ptr<FrameInfoDetectYolov8Seg> get_img();
    void enableProfile(TimeStamp *ts);
    std::string doubleToStringWithPrecision(double value, int precision);

private:
    std::shared_ptr<BMNNContext> m_bmContext;
    std::shared_ptr<BMNNNetwork> m_bmNetwork;

    std::vector<bm_image> m_resized_imgs;
    std::vector<bm_image> m_converto_imgs;

    // configuration
    float m_confThreshold = 0.25;
    float m_nmsThreshold = 0.7;

    std::vector<std::string> m_class_names;
    int m_class_num;
    int mask_len = 32;
    int m_net_h, m_net_w;
    int max_batch;
    int input_num;
    int output_num;
    int min_dim;
    int max_det = 300;
    int max_wh = 7680; // (pixels) maximum box width and height
    bmcv_convert_to_attr converto_attr;

    TimeStamp *m_ts;

    bm_handle_t handle_;
    int que_size_;
    int in_tensor_num_, out_tensor_num_;
    bm_status_t ret_;
    std::atomic<bool> stop_flag_;
    std::vector<std::shared_ptr<std::thread>> thread_vec_;

    // skip frame
    int skip_num_;
    std::vector<int> channel_skip_nums_; // channel skip numbers 的一个向量，按照channel number 存储
    // preprocess
    bm_image resized_img_;
    bmcv_convert_to_attr converto_attr_;
    std::vector<bm_image> converto_imgs_;

    // forward
    std::vector<bm_device_mem_t> input_mem_buffer_;
    int input_mem_idx_;
    std::vector<std::vector<bm_device_mem_t>> output_mem_buffer_;
    int output_mem_idx_;

    // tpu kernal post process
    float nmsThreshold_;
    float confThreshold_;
    tpu_kernel_api_yolov5NMS_t api_;
    tpu_kernel_function_t func_id;
    tpu_kernel_api_yolov8NMS_t yolov8_tpu_kernel_api_;
    bm_device_mem_t boxs_devmem_;
    bm_device_mem_t detect_num_devmem_;
    std::vector<float> boxs_sysmem_;
    int32_t detect_num_sysmem_;
    const std::vector<std::vector<std::vector<int>>> anchors{{{10, 13}, {16, 30}, {33, 23}}, {{30, 61}, {62, 45}, {59, 119}}, {{116, 90}, {156, 198}, {373, 326}}};

    // queue
    std::mutex global_frame_que_mtx_;
    /**
     * @brief 全局帧队列
     * 全局帧队列，用于存储需要预处理的图像
     */
    std::deque<std::shared_ptr<FrameInfoDetectYolov8Seg>> global_frame_que_;
    DataPipe<bm_device_mem_t *> pre_forward_que_;
    DataPipe<bm_device_mem_t *> forward_post_que_;
    // DataPipe<std::shared_ptr<FrameInfoDetectYolov8Seg>> draw_result_que_out_;
    DataPipe<std::shared_ptr<FrameInfoDetectYolov8Seg>> out_que_;

    // DataPipe<std::shared_ptr<FrameInfoDetectYolov8Seg>> get_img_que_;
    // yolov8 初始化和释放
    int init(int dev_id, std::string bmodel_file, std::string tpu_kernel_module_path, std::string mask_bmodel_path);
    int release();

    int get_input_mem_id()
    {
        return (++input_mem_idx_) % (que_size_ + 2);
    }
    int get_output_mem_id()
    {
        return (++output_mem_idx_) % (que_size_ + 2);
    }
    // yolov8 推理
    // 预处理、 推理、 后处理
    int preprocess(std::shared_ptr<bm_image> img);
    int batch_size();
    void forward_thread_dowork();
    void post_thread_dowork();
    void draw_result(cv::Mat &img, YoloV8BoxVec &result);

    void get_mask(const cv::Mat &mask_info, const cv::Mat &mask_data, const ImageInfo &para, cv::Rect bound, cv::Mat &mast_out, YoloV8Box &yolobox);
    // std::pair<double, double> get_mask_angle_and_length(const cv::Mat &mask);
    void get_mask_angle_and_length(const cv::Mat &mask, YoloV8Box &yolobox);
    std::pair<cv::Point, cv::Point> findFurthestPair(const std::vector<cv::Point> &points);
    void clip_boxes(YoloV8BoxVec &yolobox_vec, int src_w, int src_h);
    void NMS(YoloV8BoxVec &dets, float nmsConfidence);

    // tpumask model bmrt
private:
    bm_handle_t tpu_mask_handle;
    void *bmrt = nullptr;
    const bm_net_info_t *netinfo;
    std::vector<std::string> network_names;
    bool tpu_post = false;
    int tpu_mask_num = 32;
    int m_tpumask_net_h, m_tpumask_net_w;

public:
    void tpumask_Init(std::string bmodel_file, int dev_id = 0);
    void getmask_tpu(YoloV8BoxVec &boxes, int start, const bm_tensor_t &input_tensor1, Paras &paras, YoloV8BoxVec &yolobox_vec_tmp);
};

#endif